package geppetto.cat.programs;


import geppetto.cat.alignments.Alignment;
import geppetto.cat.alignments.AlignmentsSet;
import geppetto.cat.alignments.output.AlignerOutputLatex;
import geppetto.cat.models.AbstractModel;
import geppetto.cat.models.AgreementHMM;
import geppetto.cat.models.HMM;

import java.io.IOException;
import java.io.PrintStream;
import java.io.UnsupportedEncodingException;



//Compares the posteriors between agreement and baseline at decoding time

public class ComparePosteriors extends BaseExp {

	String corpusDescription;

	int size, maxSentenceSize;

	String modelDir;

	AlignmentsSet baseline, agreement;

	AbstractModel mhmmBase, mhmmAgree, mhmmBaseB, mhmmAgreeB;

	int maxSentenceOutputLen;

	String printStream;

	public ComparePosteriors(String[] args) {
		corpusDescription = args[0];
		size = Integer.parseInt(args[1]);
		maxSentenceSize = Integer.parseInt(args[2]);
		modelDir = args[3];
		maxSentenceOutputLen = Integer.parseInt(args[4]);
		printStream = args[5];
		System.out.println("Corpus " + corpusDescription);
		System.out.println("Size " + size);
		System.out.println("ModelDir " + modelDir);
		System.out.println("---");
	}

	public void initializeModels() {
		System.out.println("Loading Models");
		String dir = modelDir + "/baseline/" + forwardCorpus.getName() + "/"
				+ size + "/model/";
		mhmmBase = HMM.loadModel(forwardCorpus, dir + "/MHMM/forward");
		mhmmBaseB = HMM.loadModel(backwardCorpus, dir + "/MHMM/backward");
		dir = modelDir + "/agreement/" + forwardCorpus.getName() + "/" + size
				+ "/model/";
		mhmmAgree = AgreementHMM.loadModel(forwardCorpus, backwardCorpus, dir
				+ "/MHMM/");
		mhmmAgreeB = ((AgreementHMM) mhmmAgree).getBackwardModel();
		System.out.println("Finished Loading Model");
	}

	public void run() throws IOException {
		initializeCorpus(corpusDescription, size, maxSentenceSize);
		initializeModels();
		PrintStream out = new PrintStream(printStream);
		printBegin(out);
		printPosteriors(out);
		printEnd(out);

	}

	public void printPosterior(PrintStream out, Alignment base,
			Alignment agree, Alignment projected, Alignment gold)
			throws UnsupportedEncodingException {
		AlignerOutputLatex.outputPostLatex(base, agree, projected, gold,
				forwardCorpus, out,
				"Top Baseline, Middle Agreement, Bottom projected");
	}

	public void printPosteriorBoth(PrintStream out, Alignment base,
			Alignment baseb, Alignment agree, Alignment agreeb,
			Alignment projected, Alignment projectedb, Alignment gold,
			Alignment goldb) throws UnsupportedEncodingException {
		AlignerOutputLatex.outputPostLatexBoth(base, baseb, agree, agreeb,
				projected, projectedb, gold, goldb, forwardCorpus, out, "Sent "
						+ base.getSentenceNumber()
						+ " Top Baseline, Middle Agreement, Bottom projected");
	}

	public void printPosteriors(PrintStream out)
			throws UnsupportedEncodingException {
		//TODO There is no projected alignments for now.
		/*
		AlignmentsSet base = mhmmBase
				.viterbiAlignments(BilingualCorpus.TEST_CORPUS);
		AlignmentsSet agree = mhmmAgree
				.viterbiAlignments(BilingualCorpus.TEST_CORPUS);
		AlignmentsSet gold = forwardCorpus.getGold();
		AlignmentsSet projected = ((AgreementHMM) mhmmAgree)
				.viterbiAlignmentsWithProjection(BilingualCorpus.TEST_CORPUS);
		AlignmentsSet baseB = mhmmBaseB
				.viterbiAlignments(BilingualCorpus.TEST_CORPUS);
		AlignmentsSet agreeB = mhmmAgreeB
				.viterbiAlignments(BilingualCorpus.TEST_CORPUS);
		AlignmentsSet goldB = backwardCorpus.getGold();
		AlignmentsSet projectedB = ((AgreementHMM) mhmmAgreeB)
				.viterbiAlignmentsWithProjection(BilingualCorpus.TEST_CORPUS);
		int size = forwardCorpus.getNumSentences(BilingualCorpus.TEST_CORPUS);
		for (int i = 0; i < size; i++) {
			int fLen = forwardCorpus.getForeignSentenceLength(i,
					BilingualCorpus.TEST_CORPUS);
			int sLen = forwardCorpus.getSourceSentenceLength(i,
					BilingualCorpus.TEST_CORPUS);
			if (fLen < 12 && sLen < 12) {
				printPosteriorBoth(out, base.get(i), baseB.get(i),
						agree.get(i), agreeB.get(i), projected.get(i),
						projectedB.get(i), gold.get(i), goldB.get(i));
				out.println("\\clearpage\\pagebreak");
			} else if (fLen < maxSentenceOutputLen
					&& sLen < maxSentenceOutputLen) {
				printPosterior(out, base.get(i), agree.get(i),
						projected.get(i), gold.get(i));
				out.println("\\clearpage\\pagebreak");
			}
		}
		out.println(AlignmentEvaluator.evaluate(base, gold) + "\\\\");
		out.println(AlignmentEvaluator.evaluate(agree, gold) + "\\\\");
		out.println(AlignmentEvaluator.evaluate(projected, gold) + "\\\\");
		out.println(AlignmentEvaluator.evaluate(baseB, goldB) + "\\\\");
		out.println(AlignmentEvaluator.evaluate(agreeB, goldB) + "\\\\");
		out.println(AlignmentEvaluator.evaluate(projectedB, goldB) + "\\\\");
	*/
	}

	public static void main(String[] args) throws IOException {
		ComparePosteriors exp = new ComparePosteriors(args);
		exp.run();
	}

}
